{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9e3f209-1b47-41aa-bb33-d0e7b564203c",
   "metadata": {
    "height": 30
   },
   "source": [
    "# 部署知识库助手\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2896c72f-1aa0-4b93-aea6-45908a6e42a1",
   "metadata": {},
   "source": [
    "我们对知识库和LLM已经有了基本的理解，现在是将它们巧妙地融合并打造成一个富有视觉效果的界面的时候了。这样的界面不仅对操作更加便捷，还能便于与他人分享。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c975542-100a-431f-bfdb-2e948fd1e360",
   "metadata": {
    "height": 30
   },
   "source": [
    "Streamlit 是一种快速便捷的方法，可以直接在 **Python 中通过友好的 Web 界面演示机器学习模型**。在本课程中，我们将学习*如何使用它为生成式人工智能应用程序构建用户界面*。在构建了机器学习模型后，如果你想构建一个 demo 给其他人看，也许是为了获得反馈并推动系统的改进，或者只是因为你觉得这个系统很酷，所以想演示一下：Streamlit 可以让您通过 Python 接口程序快速实现这一目标，而无需编写任何前端、网页或 JavaScript 代码。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa74cbe-96ed-4652-a761-8740615597ed",
   "metadata": {
    "height": 30
   },
   "source": [
    "## 一、Streamlit 简介"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcf25020-22a5-435d-925a-26cbe71a5f59",
   "metadata": {},
   "source": [
    "\n",
    "`Streamlit` 是一个用于快速创建数据应用程序的开源 Python 库。它的设计目标是让数据科学家能够轻松地将数据分析和机器学习模型转化为具有交互性的 Web 应用程序，而无需深入了解 Web 开发。和常规 Web 框架，如 Flask/Django 的不同之处在于，它不需要你去编写任何客户端代码（HTML/CSS/JS），只需要编写普通的 Python 模块，就可以在很短的时间内创建美观并具备高度交互性的界面，从而快速生成数据分析或者机器学习的结果；另一方面，和那些只能通过拖拽生成的工具也不同的是，你仍然具有对代码的完整控制权。\n",
    "\n",
    "Streamlit 提供了一组简单而强大的基础模块，用于构建数据应用程序：\n",
    "\n",
    "- st.write()：这是最基本的模块之一，用于在应用程序中呈现文本、图像、表格等内容。\n",
    "\n",
    "- st.title()、st.header()、st.subheader()：这些模块用于添加标题、子标题和分组标题，以组织应用程序的布局。\n",
    "\n",
    "- st.text()、st.markdown()：用于添加文本内容，支持 Markdown 语法。\n",
    "\n",
    "- st.image()：用于添加图像到应用程序中。\n",
    "\n",
    "- st.dataframe()：用于呈现 Pandas 数据框。\n",
    "\n",
    "- st.table()：用于呈现简单的数据表格。\n",
    "\n",
    "- st.pyplot()、st.altair_chart()、st.plotly_chart()：用于呈现 Matplotlib、Altair 或 Plotly 绘制的图表。\n",
    "\n",
    "- st.selectbox()、st.multiselect()、st.slider()、st.text_input()：用于添加交互式小部件，允许用户在应用程序中进行选择、输入或滑动操作。\n",
    "\n",
    "- st.button()、st.checkbox()、st.radio()：用于添加按钮、复选框和单选按钮，以触发特定的操作。\n",
    "\n",
    "这些基础模块使得通过 Streamlit 能够轻松地构建交互式数据应用程序，并且在使用时可以根据需要进行组合和定制，更多内容请查看[官方文档](https://docs.streamlit.io/get-started)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c3209b0",
   "metadata": {},
   "source": [
    "## 二、构建应用程序"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "be6cf145-06f0-40de-98cd-2e501c3377eb",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'streamlit'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mModuleNotFoundError\u001b[39m                       Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mstreamlit\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mst\u001b[39;00m\n\u001b[32m      2\u001b[39m \u001b[38;5;66;03m# from langchain_openai import ChatOpenAI\u001b[39;00m\n\u001b[32m      3\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mos\u001b[39;00m\n",
      "\u001b[31mModuleNotFoundError\u001b[39m: No module named 'streamlit'"
     ]
    }
   ],
   "source": [
    "import streamlit as st\n",
    "# from langchain_openai import ChatOpenAI\n",
    "import os\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "import sys\n",
    "from langchain.vectorstores.chroma import Chroma\n",
    "from langchain_community.vectorstores import Milvus\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv())    # read local .env file\n",
    "from langchain_community.llms import Ollama\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "\n",
    "\n",
    "def get_llm():\n",
    "    return Ollama(base_url='http://localhost:11434', model='deepseek-r1:14b', temperature=0.1)\n",
    "\n",
    "def get_emd():\n",
    "    return OllamaEmbeddings(base_url='http://localhost:11434', model=\"bge-m3:latest\")\n",
    "\n",
    "# 初始化 Milvus 向量数据库\n",
    "def get_vectordb():\n",
    "    my_emb = get_emd()\n",
    "    # Milvus 连接参数\n",
    "    vectordb = Milvus(\n",
    "        embedding_function=my_emb,\n",
    "        collection_name=\"Vmaxs\",  # Milvus 集合名称\n",
    "        connection_args={\n",
    "            \"host\": \"192.168.0.188\",  # Milvus 服务器地址\n",
    "            \"port\": \"19530\",  # Milvus 默认端口\n",
    "        },\n",
    "    )\n",
    "    return vectordb\n",
    "\n",
    "# 不带知识库的回答\n",
    "def generate_response(input_text):\n",
    "    my_llm = get_llm()\n",
    "    output = my_llm.invoke(input_text)\n",
    "    output_parser = StrOutputParser()\n",
    "    output = output_parser.invoke(output)\n",
    "    return output\n",
    "\n",
    "# 基于知识库的问答链\n",
    "def generate_response_with_rag(question:str):\n",
    "    vectordb = get_vectordb()\n",
    "    my_llm = get_llm()\n",
    "    template = \"\"\"你是VMAX运维助手，使用以下上下文来回答问题。如果你不知道答案，就说你不知道，不要试图编造答\n",
    "    案。总是在回答的最后说“谢谢你的提问！”。\n",
    "    {context}\n",
    "    问题: {question}\n",
    "    \"\"\"\n",
    "    QA_CHAIN_PROMPT = PromptTemplate(input_variables=[\"context\",\"question\"],\n",
    "                                 template=template)\n",
    "    qa_chain = RetrievalQA.from_chain_type(my_llm,\n",
    "                                       retriever=vectordb.as_retriever(),\n",
    "                                       return_source_documents=True,\n",
    "                                       chain_type_kwargs={\"prompt\":QA_CHAIN_PROMPT})\n",
    "    result = qa_chain({\"query\": question})\n",
    "    return result[\"result\"]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def generate_response_with_rag_memory(question: str):\n",
    "    # 初始化向量数据库和LLM\n",
    "    vectordb = get_vectordb()\n",
    "    my_llm = get_llm()\n",
    "\n",
    "    memory = ConversationBufferMemory( memory_key=\"chat_history\",  # 与 prompt 的输入变量保持一致。\n",
    "    return_messages=True  # 将以消息列表的形式返回聊天记录，而不是单个字符串\n",
    "    )\n",
    "    \n",
    "    # 修改后的Prompt模板（添加chat_history变量）\n",
    "    template = \"\"\"你是VMAX运维助手，请参考以下对话历史和上下文来回答问题：\n",
    "    {chat_history}\n",
    "    \n",
    "    相关上下文：\n",
    "    {context}\n",
    "    \n",
    "    问题：{question}\n",
    "    回答结束时说“谢谢你的提问！”\n",
    "    \"\"\"\n",
    "    \n",
    "    QA_PROMPT = PromptTemplate(\n",
    "        input_variables=[\"chat_history\", \"context\", \"question\"],\n",
    "        template=template\n",
    "    )\n",
    "    \n",
    "    # 创建对话链\n",
    "    qa_chain = ConversationalRetrievalChain.from_llm(\n",
    "        llm=my_llm,\n",
    "        retriever=vectordb.as_retriever(),\n",
    "        memory=memory,\n",
    "        combine_docs_chain_kwargs={\"prompt\": QA_PROMPT},\n",
    "        chain_type=\"stuff\"\n",
    "    )\n",
    "    \n",
    "    result = qa_chain({\"question\": question})\n",
    "    return result[\"answer\"]\n",
    "\n",
    "\n",
    "# Streamlit 应用程序界面\n",
    "def main():\n",
    "    st.title('🦜🔗 VMAX-S运维助手Demo')\n",
    "    # zhipuai_api_key = st.sidebar.text_input('GLM API Key', type='password')\n",
    "\n",
    "    # 添加一个选择按钮来选择不同的模型\n",
    "    #selected_method = st.sidebar.selectbox(\"选择模式\", [\"qa_chain\", \"chat_qa_chain\", \"None\"])\n",
    "    selected_method = st.radio(\n",
    "        \"你想选择哪种模式进行对话？\",\n",
    "        [\"No-RAG\", \"generate_response_with_rag\", \"generate_response_with_rag_memory\"],\n",
    "        captions = [\"不使用基于知识库的检索问答模式\", \"基于知识库的检索问答模式\", \"基于知识库的检索问答模式（带记忆）\"])\n",
    "\n",
    "    # 用于跟踪对话历史\n",
    "    if 'messages' not in st.session_state:\n",
    "        st.session_state.messages = []\n",
    "\n",
    "    messages = st.container(height=300)\n",
    "    if prompt := st.chat_input(\"Say something\"):\n",
    "        # 将用户输入添加到对话历史中\n",
    "        st.session_state.messages.append({\"role\": \"user\", \"text\": prompt})\n",
    "\n",
    "        if selected_method == \"No-RAG\":\n",
    "            # 调用 respond 函数获取回答\n",
    "            answer = generate_response(prompt)\n",
    "        elif selected_method == \"generate_response_with_rag\":\n",
    "            answer = generate_response_with_rag(prompt)\n",
    "        elif selected_method == \"generate_response_with_rag_memory\":\n",
    "            answer = generate_response_with_rag(prompt)\n",
    "\n",
    "        # 检查回答是否为 None\n",
    "        if answer is not None:\n",
    "            # 将LLM的回答添加到对话历史中\n",
    "            st.session_state.messages.append({\"role\": \"assistant\", \"text\": answer})\n",
    "\n",
    "        # 显示整个对话历史\n",
    "        for message in st.session_state.messages:\n",
    "            if message[\"role\"] == \"user\":\n",
    "                messages.chat_message(\"user\").write(message[\"text\"])\n",
    "            elif message[\"role\"] == \"assistant\":\n",
    "                messages.chat_message(\"assistant\").write(message[\"text\"])   \n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
